#include "StdAfx.h"

#include <dshow.h>

#include "Camera.h"

#define REGISTER_FILTERGRAPH

#define SAFE_RELEASE(x) { if (x) x->Release(); x = NULL; }

enum PLAYSTATE {Stopped, Paused, Running, Init};


// Application-defined message to notify app of filtergraph events
#define WM_GRAPHNOTIFY  WM_APP+1

DWORD g_dwGraphRegister=0;

bool  g_bOneShot = false;


HRESULT WriteBitmap(PCSTR pszFileName, BITMAPINFOHEADER *pBMI, size_t cbBMI,
                    BYTE *pData, size_t cbData);



class CSampleGrabberCB : public ISampleGrabberCB 
{
public:
    long lWidth;
    long lHeight;
//    TCHAR   m_strFile[MAX_PATH];
    AM_MEDIA_TYPE *m_pMt;
    CAPTURE_COMPLETE m_pCallback;
    void*   m_pPayload;

    Image_Cap  m_Cap;

/*

    void    SetFileName(const char* strFile)
    {
        strcpy_s(m_strFile, MAX_PATH, strFile);
    }
*/
    void    SetMT(AM_MEDIA_TYPE* pMt)
    {
        m_pMt = pMt;
    }
    void    SetCallback(CAPTURE_COMPLETE pCallback)
    {
        m_pCallback = pCallback;
    }
    void    SetPlayload(void* pPayload)
    {
        m_pPayload = pPayload;
    }

    CSampleGrabberCB( ){
    } 
    STDMETHODIMP_(ULONG) AddRef() { return 2; }
    STDMETHODIMP_(ULONG) Release() { return 1; }
    STDMETHODIMP QueryInterface(REFIID riid, void ** ppv){
        if( riid == IID_ISampleGrabberCB || riid == IID_IUnknown ){ 
            *ppv = (void *) static_cast<ISampleGrabberCB*> ( this );
            return NOERROR;
        } 
        return E_NOINTERFACE;
    }
    STDMETHODIMP SampleCB( double SampleTime, IMediaSample * pSample ){
        return 0;
    }
    STDMETHODIMP BufferCB( double dblSampleTime, BYTE * pBuffer, long lBufferSize ){

        if( false == g_bOneShot)
            return 0;

        g_bOneShot= false;

        if (!pBuffer)return E_POINTER;


        // Examine the format block.
        if ((m_pMt->formattype == FORMAT_VideoInfo) && 
            (m_pMt->cbFormat >= sizeof(VIDEOINFOHEADER)) &&
            (m_pMt->pbFormat != NULL)) 
        {
            VIDEOINFOHEADER *pVih = (VIDEOINFOHEADER*)m_pMt->pbFormat;
/*
            WriteBitmap(m_strFile, &pVih->bmiHeader, 
                m_pMt->cbFormat - SIZE_PREHEADER, pBuffer, lBufferSize);
*/
            m_Cap.width = pVih->bmiHeader.biWidth;
            m_Cap.height = pVih->bmiHeader.biHeight;
            m_Cap.pBuf = pBuffer;
            m_Cap.nBufLen = lBufferSize ;
            m_Cap.type = 0;


            if( m_pCallback != NULL)
            {
                (*m_pCallback)(&m_Cap, m_pPayload);
                m_pCallback = NULL;
            }

        }
        else 
        {
            // Invalid format.
            return VFW_E_INVALIDMEDIATYPE; 
        }

        return 0;
    }

};


CSampleGrabberCB g_SGB;



#ifdef REGISTER_FILTERGRAPH

HRESULT AddGraphToRot(IUnknown *pUnkGraph, DWORD *pdwRegister) 
{
    IMoniker * pMoniker;
    IRunningObjectTable *pROT;
    WCHAR wsz[128];
    HRESULT hr;

    if (!pUnkGraph || !pdwRegister)
        return E_POINTER;

    if (FAILED(GetRunningObjectTable(0, &pROT)))
        return E_FAIL;

    hr = StringCchPrintfW(wsz, NUMELMS(wsz), L"FilterGraph %08x pid %08x\0", (DWORD_PTR)pUnkGraph, 
        GetCurrentProcessId());

    hr = CreateItemMoniker(L"!", wsz, &pMoniker);
    if (SUCCEEDED(hr)) 
    {
        // Use the ROTFLAGS_REGISTRATIONKEEPSALIVE to ensure a strong reference
        // to the object.  Using this flag will cause the object to remain
        // registered until it is explicitly revoked with the Revoke() method.
        //
        // Not using this flag means that if GraphEdit remotely connects
        // to this graph and then GraphEdit exits, this object registration 
        // will be deleted, causing future attempts by GraphEdit to fail until
        // this application is restarted or until the graph is registered again.
        hr = pROT->Register(ROTFLAGS_REGISTRATIONKEEPSALIVE, pUnkGraph, 
            pMoniker, pdwRegister);
        pMoniker->Release();
    }

    pROT->Release();
    return hr;
}


// Removes a filter graph from the Running Object Table
void RemoveGraphFromRot(DWORD pdwRegister)
{
    IRunningObjectTable *pROT;

    if (SUCCEEDED(GetRunningObjectTable(0, &pROT))) 
    {
        pROT->Revoke(pdwRegister);
        pROT->Release();
    }
}

#endif




// Writes a bitmap file
//  pszFileName:  Output file name.
//  pBMI:         Bitmap format information (including pallete).
//  cbBMI:        Size of the BITMAPINFOHEADER, including palette, if present.
//  pData:        Pointer to the bitmap bits.
//  cbData        Size of the bitmap, in bytes.

HRESULT WriteBitmap(PCSTR pszFileName, BITMAPINFOHEADER *pBMI, size_t cbBMI,
                    BYTE *pData, size_t cbData)
{
    HANDLE hFile = CreateFile(pszFileName, GENERIC_WRITE, 0, NULL, 
        CREATE_ALWAYS, 0, NULL);
    if (hFile == NULL)
    {
        return HRESULT_FROM_WIN32(GetLastError());
    }

    BITMAPFILEHEADER bmf = { };

    bmf.bfType = 'MB';
    bmf.bfSize = cbBMI+ cbData + sizeof(bmf); 
    bmf.bfOffBits = sizeof(bmf) + cbBMI; 

    DWORD cbWritten = 0;
    BOOL result = WriteFile(hFile, &bmf, sizeof(bmf), &cbWritten, NULL);
    if (result)
    {
        result = WriteFile(hFile, pBMI, cbBMI, &cbWritten, NULL);
    }
    if (result)
    {
        result = WriteFile(hFile, pData, cbData, &cbWritten, NULL);
    }

    HRESULT hr = result ? S_OK : HRESULT_FROM_WIN32(GetLastError());

    CloseHandle(hFile);

    return hr;
}

template <class T> void SafeRelease(T **ppT)
{
    if (*ppT)
    {
        (*ppT)->Release();
        *ppT = NULL;
    }
}


// Query whether a pin is connected to another pin.
//
// Note: This function does not return a pointer to the connected pin.



HRESULT IsPinConnected(IPin *pPin, BOOL *pResult)
{
    IPin *pTmp = NULL;
    HRESULT hr = pPin->ConnectedTo(&pTmp);
    if (SUCCEEDED(hr))
    {
        *pResult = TRUE;
    }
    else if (hr == VFW_E_NOT_CONNECTED)
    {
        // The pin is not connected. This is not an error for our purposes.
        *pResult = FALSE;
        hr = S_OK;
    }

    SafeRelease(&pTmp);
    return hr;
}

// Query whether a pin has a specified direction (input / output)
HRESULT IsPinDirection(IPin *pPin, PIN_DIRECTION dir, BOOL *pResult)
{
    PIN_DIRECTION pinDir;
    HRESULT hr = pPin->QueryDirection(&pinDir);
    if (SUCCEEDED(hr))
    {
        *pResult = (pinDir == dir);
    }
    return hr;
}


// Match a pin by pin direction and connection state.
HRESULT MatchPin(IPin *pPin, PIN_DIRECTION direction, BOOL bShouldBeConnected, BOOL *pResult)
{
    ATLASSERT(pResult != NULL);

    BOOL bMatch = FALSE;
    BOOL bIsConnected = FALSE;

    HRESULT hr = IsPinConnected(pPin, &bIsConnected);
    if (SUCCEEDED(hr))
    {
        if (bIsConnected == bShouldBeConnected)
        {
            hr = IsPinDirection(pPin, direction, &bMatch);
        }
    }

    if (SUCCEEDED(hr))
    {
        *pResult = bMatch;
    }
    return hr;
}


// Return the first unconnected input pin or output pin.
HRESULT FindUnconnectedPin(IBaseFilter *pFilter, PIN_DIRECTION PinDir, IPin **ppPin)
{
    IEnumPins *pEnum = NULL;
    IPin *pPin = NULL;
    BOOL bFound = FALSE;

    HRESULT hr = pFilter->EnumPins(&pEnum);
    if (FAILED(hr))
    {
        goto done;
    }

    while (S_OK == pEnum->Next(1, &pPin, NULL))
    {
        hr = MatchPin(pPin, PinDir, FALSE, &bFound);
        if (FAILED(hr))
        {
            goto done;
        }
        if (bFound)
        {
            *ppPin = pPin;
            (*ppPin)->AddRef();
            break;
        }
        SafeRelease(&pPin);
    }

    if (!bFound)
    {
        hr = VFW_E_NOT_FOUND;
    }

done:
    SafeRelease(&pPin);
    SafeRelease(&pEnum);
    return hr;
}


HRESULT ConnectFilters(
                       IGraphBuilder *pGraph, // Filter Graph Manager.
                       IPin *pOut,            // Output pin on the upstream filter.
                       IBaseFilter *pDest)    // Downstream filter.
{
    IPin *pIn = NULL;

    // Find an input pin on the downstream filter.
    HRESULT hr = FindUnconnectedPin(pDest, PINDIR_INPUT, &pIn);
    if (SUCCEEDED(hr))
    {
        // Try to connect them.
        hr = pGraph->Connect(pOut, pIn);
        pIn->Release();
    }
    return hr;
}


// Connect filter to filter

HRESULT ConnectFilters(IGraphBuilder *pGraph, IBaseFilter *pSrc, IBaseFilter *pDest)
{
    IPin *pOut = NULL;

    // Find an output pin on the first filter.
    HRESULT hr = FindUnconnectedPin(pSrc, PINDIR_OUTPUT, &pOut);
    if (SUCCEEDED(hr))
    {
        hr = ConnectFilters(pGraph, pOut, pDest);
        pOut->Release();
    }
    return hr;
}



PLAYSTATE g_psCurrent;

CCamera::CCamera(void)
{
    g_pVW = NULL;
    g_pMC = NULL;
    g_pME = NULL;
    g_pGraph = NULL;
    g_pCapture = NULL;
    g_psCurrent = Stopped;

    pGrabberF = NULL;
    pGrabber = NULL;
    pNullF = NULL;
}

CCamera::~CCamera(void)
{

}

HRESULT CCamera::InitFilters(HWND hWnd)
{
    HRESULT hr;


    m_hWnd = hWnd;

    // Create the filter graph
    hr = CoCreateInstance (CLSID_FilterGraph, NULL, CLSCTX_INPROC,
        IID_IGraphBuilder, (void **) &g_pGraph);
    if (FAILED(hr))
        return hr;

    // Create the capture graph builder
    hr = CoCreateInstance (CLSID_CaptureGraphBuilder2 , NULL, CLSCTX_INPROC,
        IID_ICaptureGraphBuilder2, (void **) &g_pCapture);
    if (FAILED(hr))
        return hr;

    // Obtain interfaces for media control and Video Window
    hr = g_pGraph->QueryInterface(IID_IMediaControl,(LPVOID *) &g_pMC);
    if (FAILED(hr))
        return hr;

    hr = g_pGraph->QueryInterface(IID_IVideoWindow, (LPVOID *) &g_pVW);
    if (FAILED(hr))
        return hr;

    hr = g_pGraph->QueryInterface(IID_IMediaEventEx, (LPVOID *) &g_pME);
    if (FAILED(hr))
        return hr;

    // Set the window handle used to process graph events
    hr = g_pME->SetNotifyWindow((OAHWND)m_hWnd, WM_GRAPHNOTIFY, 0);
    if (FAILED(hr))
    {
        ATLTRACE(TEXT("Failed to get video interfaces!  hr=0x%x"), hr);
        return hr;
    }

    // Attach the filter graph to the capture graph
    hr = g_pCapture->SetFiltergraph(g_pGraph);
    if (FAILED(hr))
    {
        ATLTRACE(TEXT("Failed to set capture filter graph!  hr=0x%x"), hr);
        return hr;
    }

    hr = FindCaptureDevice(&pSrcFilter);
    if (FAILED(hr))
    {
        // Don't display a message because FindCaptureDevice will handle it
        return hr;
    }
    // Add Capture filter to our graph.
    hr = g_pGraph->AddFilter(pSrcFilter, L"Video Capture");
    if (FAILED(hr))
    {
        ATLTRACE(TEXT("Couldn't add the capture filter to the graph!  hr=0x%x\r\n\r\n") 
            TEXT("If you have a working video capture device, please make sure\r\n")
            TEXT("that it is connected and is not being used by another application.\r\n\r\n")
            TEXT("The sample will now close."), hr);
        pSrcFilter->Release();
        return hr;
    }


#ifdef REGISTER_FILTERGRAPH
    // Add our graph to the running object table, which will allow
    // the GraphEdit application to "spy" on our graph
    hr = AddGraphToRot(g_pGraph, &g_dwGraphRegister);
    if (FAILED(hr))
    {
        ATLTRACE(TEXT("Failed to register filter graph with ROT!  hr=0x%x"), hr);
        g_dwGraphRegister = 0;
    }
#endif


    return hr;

}

bool CCamera::IsRun()
{
    return g_psCurrent == Running;
}

HRESULT CCamera::Init(HWND hWnd)
{
    HRESULT hr = S_OK;
    hr =InitFilters(hWnd);
    if( FAILED(hr))
    {
        return hr;
    }

    // Create the Sample Grabber filter.
    hr = CoCreateInstance(CLSID_SampleGrabber, NULL, CLSCTX_INPROC_SERVER,
        IID_PPV_ARGS(&pGrabberF));
    if (FAILED(hr))
    {
        return hr;
    }

    hr = g_pGraph->AddFilter(pGrabberF, L"Sample Grabber");
    if (FAILED(hr))
    {
        return hr;
    }

    hr = pGrabberF->QueryInterface(IID_PPV_ARGS(&pGrabber));
    if (FAILED(hr))
    {
        return hr;
    }


    ZeroMemory(&mt, sizeof(mt));
    mt.majortype = MEDIATYPE_Video;
    mt.subtype = MEDIASUBTYPE_RGB24;

    hr = pGrabber->SetMediaType(&mt);
    if (FAILED(hr))
    {
        return hr;
    }



    hr = ConnectFilters(g_pGraph, pSrcFilter, pGrabberF);
    if (FAILED(hr))
    {
        return hr;
    }

    IPin *pOut = NULL;

    // Find an output pin on the first filter.
    hr = FindUnconnectedPin(pGrabberF, PINDIR_OUTPUT, &pOut);

    if (SUCCEEDED(hr))
    {
        hr = g_pGraph->Render(pOut);
        if( FAILED(hr))
            return hr;
        pOut->Release();
    }

    // Now that the filter has been added to the graph and we have
    // rendered its stream, we can release this reference to the filter.
    pSrcFilter->Release();

    // Set video window style and position

    // Set the video window to be a child of the main window
    hr = g_pVW->put_Owner((OAHWND)m_hWnd);
    if (FAILED(hr))
        return hr;

    // Set video window style
    hr = g_pVW->put_WindowStyle(WS_CHILD | WS_CLIPCHILDREN);
    if (FAILED(hr))
        return hr;

    // Use helper function to position video window in client rect 
    // of main application window
    if (g_pVW)
    {
        RECT rc;

        // Make the preview video fill our window
        GetClientRect(m_hWnd, &rc);
        g_pVW->SetWindowPosition(0, 0, rc.right, rc.bottom);
    }

    // Make the video window visible, now that it is properly positioned
    hr = g_pVW->put_Visible(OATRUE);
    if (FAILED(hr))
        return hr;

    hr = g_pVW->put_WindowState(WM_SHOWWINDOW);


    if (FAILED(hr))
    {
        ATLTRACE(TEXT("Couldn't initialize video window!  hr=0x%x"), hr);
        return hr;
    }


    return S_OK;
}

void CCamera::Release()
{
#ifdef REGISTER_FILTERGRAPH
    // Remove filter graph from the running object table   
    if (g_dwGraphRegister)
        RemoveGraphFromRot(g_dwGraphRegister);
#endif

    // Release DirectShow interfaces
    SAFE_RELEASE(pGrabberF);
    SAFE_RELEASE(pGrabber);
    SAFE_RELEASE(pNullF);
    SAFE_RELEASE(g_pMC);
    SAFE_RELEASE(g_pME);
    SAFE_RELEASE(g_pVW);
    SAFE_RELEASE(g_pGraph);
    SAFE_RELEASE(g_pCapture);
}

HRESULT CCamera::start()
{
    // Start previewing video data
    HRESULT hr;
    hr = g_pMC->Run();
    if (FAILED(hr))
    {
        ATLTRACE(TEXT("Couldn't run the graph!  hr=0x%x"), hr);
        return hr;
    }

    // Remember current state
    g_psCurrent = Running;

    return S_OK;
}

HRESULT CCamera::CaptureBuf( CAPTURE_COMPLETE callback, void* payload)
{
    HRESULT hr;
    BYTE *pBuffer = NULL;

    hr = pGrabber->SetOneShot(FALSE);
    if (FAILED(hr))
    {
        return hr;
    }

    hr = pGrabber->SetBufferSamples(TRUE);
    if (FAILED(hr))
    {
        return hr;
    }

    hr = pGrabber->GetConnectedMediaType(&mt);
    if (FAILED(hr))
    {
        return hr;
    }


    g_bOneShot = true;

    g_SGB.SetMT(&mt);
    g_SGB.SetCallback(callback);
    g_SGB.SetPlayload(payload);
    hr = pGrabber->SetCallback(&g_SGB, 1);
    if (FAILED(hr))
    {
        return hr;
    }

    return hr;
}

HRESULT CCamera::CaptureFile(const char* strFile, CAPTURE_COMPLETE callback, void* payload)
{
    HRESULT hr;
    BYTE *pBuffer = NULL;

    hr = pGrabber->SetOneShot(FALSE);
    if (FAILED(hr))
    {
        return hr;
    }

    hr = pGrabber->SetBufferSamples(TRUE);
    if (FAILED(hr))
    {
        return hr;
    }

    hr = pGrabber->GetConnectedMediaType(&mt);
    if (FAILED(hr))
    {
        return hr;
    }


    g_bOneShot = true;

//    g_SGB.SetFileName(strFile);
    g_SGB.SetMT(&mt);
    g_SGB.SetCallback(callback);
    g_SGB.SetPlayload(payload);
    hr = pGrabber->SetCallback(&g_SGB, 1);
    if (FAILED(hr))
    {
        return hr;
    }

    return hr;
}


void CCamera::stop()
{
    // Stop previewing data
    if (g_pMC)
        g_pMC->StopWhenReady();

    g_psCurrent = Stopped;

    // Stop receiving events
    if (g_pME)
        g_pME->SetNotifyWindow(NULL, WM_GRAPHNOTIFY, 0);

    // Relinquish ownership (IMPORTANT!) of the video window.
    // Failing to call put_Owner can lead to assert failures within
    // the video renderer, as it still assumes that it has a valid
    // parent window.
    if(g_pVW)
    {
        g_pVW->put_Visible(OAFALSE);
        g_pVW->put_Owner(NULL);
    }



}


HRESULT CCamera::FindCaptureDevice(IBaseFilter ** ppSrcFilter)
{
    HRESULT hr = S_OK;
    IBaseFilter * pSrc = NULL;
    IMoniker* pMoniker =NULL;
    ICreateDevEnum *pDevEnum =NULL;
    IEnumMoniker *pClassEnum = NULL;

    if (!ppSrcFilter)
    {
        return E_POINTER;
    }

    // Create the system device enumerator
    hr = CoCreateInstance (CLSID_SystemDeviceEnum, NULL, CLSCTX_INPROC,
        IID_ICreateDevEnum, (void **) &pDevEnum);
    if (FAILED(hr))
    {
        ATLTRACE(TEXT("Couldn't create system enumerator!  hr=0x%x"), hr);
    }

    // Create an enumerator for the video capture devices

    if (SUCCEEDED(hr))
    {
        hr = pDevEnum->CreateClassEnumerator (CLSID_VideoInputDeviceCategory, &pClassEnum, 0);
        if (FAILED(hr))
        {
            ATLTRACE(TEXT("Couldn't create class enumerator!  hr=0x%x"), hr);
        }
    }

    if (SUCCEEDED(hr))
    {
        // If there are no enumerators for the requested type, then 
        // CreateClassEnumerator will succeed, but pClassEnum will be NULL.
        if (pClassEnum == NULL)
        {
            MessageBox(m_hWnd,TEXT("No video capture device was detected.\r\n\r\n")
                TEXT("This sample requires a video capture device, such as a USB WebCam,\r\n")
                TEXT("to be installed and working properly.  The sample will now close."),
                TEXT("No Video Capture Hardware"), MB_OK | MB_ICONINFORMATION);
            hr = E_FAIL;
        }
    }

    // Use the first video capture device on the device list.
    // Note that if the Next() call succeeds but there are no monikers,
    // it will return S_FALSE (which is not a failure).  Therefore, we
    // check that the return code is S_OK instead of using SUCCEEDED() macro.

    if (SUCCEEDED(hr))
    {
        hr = pClassEnum->Next (1, &pMoniker, NULL);
        if (hr == S_FALSE)
        {
            ATLTRACE(TEXT("Unable to access video capture device!"));   
            hr = E_FAIL;
        }
    }

    if (SUCCEEDED(hr))
    {
        // Bind Moniker to a filter object
        hr = pMoniker->BindToObject(0,0,IID_IBaseFilter, (void**)&pSrc);
        if (FAILED(hr))
        {
            ATLTRACE(TEXT("Couldn't bind moniker to filter object!  hr=0x%x"), hr);
        }
    }

    // Copy the found filter pointer to the output parameter.
    if (SUCCEEDED(hr))
    {
        *ppSrcFilter = pSrc;
        (*ppSrcFilter)->AddRef();
    }

    SAFE_RELEASE(pSrc);
    SAFE_RELEASE(pMoniker);
    SAFE_RELEASE(pDevEnum);
    SAFE_RELEASE(pClassEnum);

    return hr;
}

